"""
@Time : 2024/03/30
@Author : JiBingYu
@File : try_new_env.py
@Description : 目的是通过这个文件进行输入和输出的定义，完成在其他环境中调用此环境即可使用训练好的模型
但是，这个返回的是model，需要外部环境和env类似才可以运行，因此考虑使用新的模型调用方式，单独对每个智能体进行环境的调用。

Input: 全局环境信息【在env_core中对信息进行划分】
obs = np.array 【1，3，obs_dim】
np.array([
            [[29, 14, 0, 0, -7.75, 14, 0, 0],
            [-34, -40, 0, 0, -7.75, 14, 0, 0],
            [26, 21, 0, 0, -7.75, 14, 0, 0]]
            ])
Output:目标和友方的运动方式【在env_core中修改动作部分】
action = list 【1，3，2】
[[[-0.21827637 -0.32159883]
  [-0.12455051 -0.7943451 ]
  [-0.19297189  2.5175028 ]]]
使用方法

from train.try_new_env import Get_Run
model = Get_Run(models_dir)
LOOP:
    eval_actions_env, image = model.step_eval(obs)
    images.append(image)
    obs = env.step()

from PIL import Image
import imageio
imageio.mimsave('GGIF_dir'+'episode.gif', images, 'GIF', duration=1)
"""

# !/usr/bin/env python
import setproctitle
import numpy as np
from pathlib import Path
import torch
import sys
import os

current_dir = os.path.dirname(__file__)
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)

from config_true_env import get_config
from envs.env_wrappers import DummyVecEnv
from config_env import get_env_config
from envs.ruler import ruler_main, show_move_list
parser_ev = get_env_config()

"""Train script for MPEs."""

def make_train_env(all_args, env_size = []):
    def get_env_fn(rank):
        def init_env(env_size):
            from envs.env_continuous import ContinuousActionEnv
            env = ContinuousActionEnv(env_size)
            env.seed(all_args.seed + rank * 1000)
            return env
        return init_env
    return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)], env_size)


def make_eval_env(all_args, env_size = []):
    def get_env_fn(rank):
        def init_env(env_size):
            from envs.env_continuous import ContinuousActionEnv
            env = ContinuousActionEnv(env_size)
            env.seed(all_args.seed + rank * 1000)
            return env
        return init_env
    return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_eval_rollout_threads)], env_size)


def parse_args(args, parser):
    parser.add_argument("--scenario_name", type=str, default="MyEnv", help="Which scenario to run on")
    parser.add_argument("--num_landmarks", type=int, default=3)
    parser.add_argument("--num_agents", type=int, default=parser_ev["agent_num"], help="number of players")

    all_args = parser.parse_known_args(args)[0]

    return all_args

def Get_Run(model_dir, env_size = []):
    parser = get_config()
    all_args = parse_args(sys.argv[1:], parser)

    if all_args.algorithm_name == "rmappo":
        assert all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy, "check recurrent policy!"
    elif all_args.algorithm_name == "mappo":
        assert (
            all_args.use_recurrent_policy == False and all_args.use_naive_recurrent_policy == False
        ), "check recurrent policy!"
    else:
        raise NotImplementedError

    assert (
        all_args.share_policy == True and all_args.scenario_name == "simple_speaker_listener"
    ) == False, "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    # if not run_dir.exists():
    #     os.makedirs(str(run_dir))

    if not run_dir.exists():
        curr_run = "run1"
    else:
        exst_run_nums = [
            int(str(folder.name).split("run")[1])
            for folder in run_dir.iterdir()
            if str(folder.name).startswith("run")
        ]
        if len(exst_run_nums) == 0:
            curr_run = "run1"
        else:
            curr_run = "run%i" % (max(exst_run_nums) + 1)
    run_dir = run_dir / curr_run
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_train_env(all_args, env_size)
    eval_envs = make_eval_env(all_args, env_size) if all_args.use_eval else None
    num_agents = all_args.num_agents
    # model_dir = "D:/Reinforce/light_mappo-main/light_mappo-main/results/MyEnv/MyEnv/mappo/check/run125/models"
    all_args.model_dir = model_dir
    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from runner.shared.env_runner import EnvRunner as Runner
    else:
        from runner.separated.env_runner import EnvRunner as Runner

    runner = Runner(config)
    return runner
